# @Time : 2021/4/25 10:47
# @Author : Fioman 
# @Phone : 13149920693
import os

import cv2 as cv
import imutils
from tools.fioman_math_tools import *

"""
图像的基本操作,比如获取平均像素,寻找图像物料区域,等等.
"""


def image_show_for_operation(name, image, showState=True):
    """
    显示图片
    :param name:显示的窗口名称
    :param image:要显示的图片
    :param showState:是否打开显示
    :return:
    """
    if showState:
        scaleRate = 6
        showWidth = int(image.shape[1] / scaleRate)
        showHeight = int(image.shape[0] / scaleRate)
        showImage = cv.resize(image, (showWidth, showHeight), interpolation=cv.INTER_AREA)
        cv.namedWindow("{}".format(name))
        cv.imshow("{}".format(name), showImage)
        cv.waitKey(0)


def get_center_by_contour(cnt):
    """
    根据传入的轮廓获取到轮廓的中心
    :param cnt:传入的轮廓
    :return:因为是轮廓的中心,默认返回int类型
    """
    (centerX, centerY), *_ = cv.minAreaRect(cnt)
    return int(centerX), int(centerY)


def get_rect_points_clockwise(rectPoints):
    """
    根据矩形的四个顶点,获取其顺时针方向的四个顶点
    :param rectPoints:
    :return:
    """
    rectPointsSorted = sorted(rectPoints, key=np.sum)
    leftTop = rectPointsSorted[0]
    rightBottom = rectPointsSorted[-1]
    rightTopX = int(rectPointsSorted[1][0])
    leftBottomX = int(rectPointsSorted[2][0])
    if rightTopX > leftBottomX:
        rightTop = rectPointsSorted[1]
        leftBottom = rectPointsSorted[2]
    else:
        rightTop = rectPointsSorted[2]
        leftBottom = rectPointsSorted[1]

    return leftTop, rightTop, rightBottom, leftBottom


def delete_table(self, tableName):
    self.cursor.execute("delete from {}".format(tableName))


def find_single_image_material_area(image):
    """
    返回单个图像的物料区域位置和旋转角度.
    :param image:传入的原始图片,在不知道物料区域面积的情况下来进行寻找.
    :param imageType: 传入进来的图片是左图还是右图,还是拼接后的图,如果是左图,写字的位置就要靠近左边,如果是右图,写字的位置就要靠近右边
    :return:带物料区域框的物料图,还有就是带字的平均像素
    """
    thresUsed = np.mean(image[:150, :1000]) + 15
    filterArea = 2000  # 单位像素,大于100个像素的区域认为是有效区域
    T, board = cv.threshold(image, thresUsed, 255, cv.THRESH_BINARY)
    cnts = cv.findContours(board.copy(), cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)

    # 然后是连通所有的满足要求的区域,通过白线连接起来,形成整个物料框
    largeContours = [cnt for cnt in cnts if cv.contourArea(cnt) > filterArea]
    showContours = cv.cvtColor(image, cv.COLOR_GRAY2BGR)
    contoursSorted = sorted([cv.contourArea(cnt) for cnt in largeContours])
    print("len(largeContours) = {}, contourArea = {}".format(len(contoursSorted), contoursSorted))
    cv.drawContours(showContours, largeContours, -1, (0, 255, 0), 3)
    image_show_for_operation("ContoursShow", showContours)

    if len(largeContours) >= 2:
        centerPoints = [get_center_by_contour(contour) for contour in largeContours]
        for index, center in enumerate(centerPoints):
            anotherCenter = centerPoints[(index + 1) % len(centerPoints)]
            cv.line(board, center, anotherCenter, 255, 20)
        image_show_for_operation("JointCenterLine", board)

        # 连线之后重新再找一遍轮廓
        cnts = cv.findContours(board.copy(), cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
        cnts = imutils.grab_contours(cnts)
        cntFirstFinded = sorted(cnts, key=cv.contourArea, reverse=True)[0]
    else:
        cntFirstFinded = largeContours[0]

    # 这个轮廓找到以后,将这个轮廓的四个角去除掉
    rectPoints = np.int0(cv.boxPoints(cv.minAreaRect(cntFirstFinded)))
    leftTop, rightTop, rightBottom, leftBottom = get_rect_points_clockwise(rectPoints)
    # 去除掉物料区域的四个角,这个如何去除.移除掉整个宽度1/5,剩下五分之三进行检测
    totalWidth = distance(leftTop, rightTop)
    totalHeight = distance(leftTop, leftBottom)
    widthMoved = int(totalWidth * 0.2)
    heightMoved = int(totalHeight * 0.15)
    extend = 20
    boardFinal = image.copy()
    boardFinal[leftTop[1] - extend:leftTop[1] + heightMoved, leftTop[0] - extend:leftTop[0] + widthMoved] = 0
    boardFinal[rightTop[1] - extend:rightTop[1] + heightMoved, rightTop[0] - widthMoved:rightTop[0] + extend] = 0
    boardFinal[leftBottom[1] - heightMoved:leftBottom[1] + extend, leftBottom[0] - extend:leftBottom[0] + widthMoved] = 0
    boardFinal[rightBottom[1] - heightMoved:rightBottom[1] + extend, rightBottom[0] - widthMoved:rightBottom[0] + extend] = 0
    image_show_for_operation("FinalBoardAngleMoved", boardFinal)

    T, board = cv.threshold(boardFinal, thresUsed, 255, cv.THRESH_BINARY)
    image_show_for_operation("FinalBoardThres", board)

    # 去掉四个角以后再找一遍轮廓
    cnts = cv.findContours(board.copy(), cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)
    largeContours = [cnt for cnt in cnts if cv.contourArea(cnt) > filterArea]

    if len(largeContours) >= 2:
        centerPoints = [get_center_by_contour(contour) for contour in largeContours]
        for index, center in enumerate(centerPoints):
            anotherCenter = centerPoints[(index + 1) % len(centerPoints)]
            cv.line(board, center, anotherCenter, 255, 20)
        image_show_for_operation("JointCenterLine", board)

        # 连线之后再找一遍
        cnts = cv.findContours(board.copy(), cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
        cnts = imutils.grab_contours(cnts)
        cntFinalFinded = sorted(cnts, key=cv.contourArea, reverse=True)[0]
    else:
        cntFinalFinded = largeContours[0]

    totalBoardThres = boardFinal.copy()
    cv.drawContours(totalBoardThres, [cntFinalFinded], -1, 255, -1)
    image_show_for_operation("TotalBoardThres", totalBoardThres)

    boxMaster = cv.minAreaRect(cntFinalFinded)
    (centerX, centerY), (w, h), angle = boxMaster
    if angle > 70:
        angle = 90 - angle
        boxMaster = (centerX, centerY), (h, w), angle
    elif angle < - 70:
        angle = 90 + angle
        boxMaster = (centerX, centerY), (h, w), angle

    rectPoints = np.int0(cv.boxPoints(boxMaster))
    # 将图的信息画上去
    boxMask = np.zeros(image.shape, dtype=np.uint8)
    cv.drawContours(boxMask, [rectPoints], -1, 255, -1)
    imageBlackAround = cv.bitwise_and(image, boxMask)
    # 求出物料区域的平均像素,根据图像和物料区域的位置,返回旋转之后的物料区域图形.
    boardWithoutAngle = get_board_widthout_angle(image, boxMaster)
    meanVal = round(float(np.mean(boardWithoutAngle)), 2)

    return imageBlackAround, boxMaster, meanVal

def get_size_scale_image(imageSrc, sizeScale):
    """
    高宽同等比例缩放图像,缩放比例
    :param imageSrc: 原图像
    :param sizeScale: 缩放比例
    :return:
    """
    h, w = imageSrc.shape[:2]
    newWidth, newHeight = int(w / sizeScale), int(h / sizeScale)
    newImage = cv.resize(imageSrc, (newWidth, newHeight), cv.INTER_AREA)
    return newImage


def get_board_widthout_angle(image, boxMaster, extend=False):
    """
    根据原图和物料的位置,返回旋转之后的物料图.
    :param image:原图
    :param boxMaster:物料区域的位置.
    :return:返回旋转之后的物料图.
    """
    (centerX, centerY), (w, h), angle = boxMaster
    widthHalf = int(w / 2)
    heightHalf = int(h / 2)
    yStart = max(int(centerY - heightHalf), 0)  # 之所以有零的情况,是因为旋转之后物料框的y起点可能跑外面去了,也就是说cy-heightHalf可能变成负数了
    yEnd = min(int(centerY + heightHalf), image.shape[0])  # y也是一样的,旋转之后图跑外面去了,这个时候结束的位置应该是图像的高度
    xStart = max(int(centerX - widthHalf), 0)
    xEnd = min(int(centerX + widthHalf), image.shape[1])
    M = cv.getRotationMatrix2D((centerX, centerY), angle, 1)
    if extend:
        extend = 50
        if yStart - extend > 0:
            yStart = yStart - extend
        if yEnd + extend < image.shape[0]:
            yEnd = yEnd + extend
        if xStart - extend > 0:
            xStart = xStart - extend
        if xEnd + extend < image.shape[1]:
            xEnd = xEnd + extend

    totalBoard = cv.warpAffine(image, M, (image.shape[1], image.shape[0]))[yStart:yEnd, xStart:xEnd]
    return totalBoard


def write_board_info_to_image(image, boxMaster, writePos, writeInfo):
    """
    将位置信息写入到图上去
    :param image:要写的图像
    :param boxMaster:位置信息
    :param writePos:书写的位置
    :param writeInfo:书写的内容
    :return:返回写好信息的彩色图片
    """
    imageColor = cv.cvtColor(image, cv.COLOR_GRAY2BGR)
    singleBox = boxMaster
    rectPoints = np.int0(cv.boxPoints(singleBox))
    cv.drawContours(imageColor, [rectPoints], -1, (0, 255, 0), 3)
    if writePos == "center":
        singleInfo = writeInfo[0]
        (centerX, centerY), *_ = singleBox
        cv.putText(imageColor, singleInfo, (int(centerX) - 200, int(centerY) - 200), cv.FONT_HERSHEY_COMPLEX, 3, (0, 0, 255), 4)

    elif writePos == "leftAndRight":
        singleBox = boxMaster
        leftInfo, rightInfo = writeInfo
        leftTop, rightTop, *_ = get_rect_points_clockwise(np.int0(cv.boxPoints(singleBox)))
        cv.putText(imageColor, leftInfo, (leftTop[0] + 20, leftTop[1] + 300), cv.FONT_HERSHEY_COMPLEX, 2, (0, 0, 255), 2)
        cv.putText(imageColor, rightInfo, (rightTop[0] - 500, rightTop[1] + 500), cv.FONT_HERSHEY_COMPLEX, 2, (0, 0, 255), 2)

    elif writePos == "imageNum":
        singleInfo = writeInfo
        (centerX, centerY), *_ = singleBox
        cv.putText(imageColor, singleInfo, (int(centerX) - 200, int(centerY) - 200), cv.FONT_HERSHEY_COMPLEX, 5, (0, 0, 255), 5)
    image_show_for_operation("BoardColorWithInfo", imageColor)

    return imageColor


if __name__ == '__main__':
    filePath = r"D:\totalBoardTest"
    keepDir = filePath
    if not os.path.isdir(keepDir):
        keepDir, _ = os.path.split(filePath)

    keepOkPath = os.path.join(keepDir, "finished")
    keepResPath = os.path.join(keepDir, "total_board_res")
    keepFailPath = os.path.join(keepDir, "fail")

    if not os.path.exists(keepOkPath):
        os.mkdir(keepOkPath)
    if not os.path.exists(keepResPath):
        os.mkdir(keepResPath)
    if not os.path.exists(keepFailPath):
        os.mkdir(keepFailPath)

    writePos = "center"
    fileNames = []
    if os.path.isfile(filePath):
        filePath, fileName = os.path.split(filePath)
        fileNames.append(fileName)
    else:
        for root, dirs, files in os.walk(filePath):
            if root != filePath:
                continue
            for file in files:
                if file[-4:] != ".bmp":
                    continue
                fileNames.append(file)

    for index, fileName in enumerate(fileNames):
        print("*" * 10 + " 第 {} 张图 ({})".format(index + 1, fileName) + "*" * 10)
        filePathReal = os.path.join(filePath, fileName)
        imgSrc = cv.imread(filePathReal, cv.IMREAD_GRAYSCALE)
        boardWithColor, boxMaster, meanValue = find_single_image_material_area(imgSrc)
        writeInfo = "meanVal = {}".format(meanValue)
        boardWithInfo = write_board_info_to_image(imgSrc, boxMaster, writePos, [writeInfo])
        if boardWithInfo is not None:
            cv.imwrite(os.path.join(keepOkPath, fileName), imgSrc)
            os.remove(filePathReal)
            cv.imwrite(os.path.join(keepResPath, fileName), boardWithInfo)
        else:
            cv.imwrite(os.path.join(keepFailPath, fileName), imgSrc)
            os.remove(filePathReal)
